import torch

import myutils
import matplotlib.pyplot as plt
import Net
n_epochs = 50
batch_size_train = 256
batch_size_test = 256
learning_rate = 0.00001
random_seed = 1
torch.manual_seed(random_seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载数据集
train, test = myutils.load_data_mnist(batch_size_train, batch_size_test)

"""
example = enumerate(train)
batch_idx, (example_data, example_targets) = next(example)

fig = plt.figure()
for i in range(6):
  plt.subplot(2,3,i+1)
  plt.tight_layout()
  plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
  plt.title("Ground Truth: {}".format(example_targets[i]))
  plt.xticks([])
  plt.yticks([])
plt.show()
"""

net = Net.LeNet()
# 采用adam梯度下降
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
myutils.train_net(net, train, test, optimizer, device, n_epochs)
# 保存网络
# torch.save(net, 'net.pth')